"""
Toolkit模式示例脚本
演示如何使用功能组件架构
"""
import os
import sys
import argparse
import logging
from typing import Dict, Any, List

# 添加项目根目录到路径
sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), "../../../")))

from lib.llm.providers.enhanced_moonshot import EnhancedMoonshotProvider
from lib.llm.provider_registry import get_registry

# 设置日志
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(name)s - %(levelname)s - %(message)s')
logger = logging.getLogger("toolkit_demo")

def register_providers():
    """注册增强版提供商"""
    registry = get_registry()
    registry.register("enhanced_moonshot", EnhancedMoonshotProvider)
    logger.info("已注册增强版提供商")

def create_provider(provider_type: str, api_key: str = None, **kwargs) -> Any:
    """
    创建提供商实例
    
    Args:
        provider_type: 提供商类型
        api_key: API密钥（可选）
        **kwargs: 其他配置参数
        
    Returns:
        提供商实例
    """
    config = {
        "provider": provider_type,
        "model": kwargs.get("model", "moonshot-v1-32k"),
        "temperature": kwargs.get("temperature", 0.1),
        "max_tokens": kwargs.get("max_tokens", 2048),
        "search": kwargs.get("search", False)
    }
    
    # 设置API密钥
    if api_key:
        config["api_key"] = api_key
    else:
        # 从环境变量获取
        env_var = f"{provider_type.upper()}_API_KEY".replace("ENHANCED_", "")
        config["api_key"] = os.environ.get(env_var)
    
    registry = get_registry()
    provider_class = registry.get(provider_type)
    
    if not provider_class:
        raise ValueError(f"未知的提供商类型: {provider_type}")
        
    return provider_class(config)

def demo_partial_mode(provider: Any):
    """
    演示部分输出功能
    
    Args:
        provider: 提供商实例
    """
    logger.info("演示部分输出功能")
    
    if not provider.has_capability("partial_mode"):
        logger.error("当前提供商不支持部分输出功能")
        return
        
    partial_mode = provider.get_capability("partial_mode")
    
    # 设置初始消息上下文
    provider.set_current_messages([
        {"role": "system", "content": "你是一个有用的助手。"},
        {"role": "user", "content": "请写一段关于人工智能的短文。"}
    ])
    
    # 使用部分输出功能
    prefix = "人工智能(AI)是计算机科学的一个分支，"
    result = partial_mode.continue_text(prefix)
    
    print("\n部分输出结果:\n" + result)

def demo_web_search(provider: Any):
    """
    演示网络搜索功能
    
    Args:
        provider: 提供商实例
    """
    logger.info("演示网络搜索功能")
    
    if not provider.has_capability("web_search"):
        logger.error("当前提供商不支持网络搜索功能")
        return
    
    # 发送包含搜索查询的消息
    response = provider.complete([
        {"role": "user", "content": "最近有哪些科技新闻？"}
    ])
    
    print("\n搜索结果:\n" + response)

def demo_tool_calls(provider: Any):
    """
    演示工具调用功能
    
    Args:
        provider: 提供商实例
    """
    logger.info("演示工具调用功能")
    
    if not provider.has_capability("tool_calls"):
        logger.error("当前提供商不支持工具调用功能")
        return
        
    tool_call = provider.get_capability("tool_calls")
    
    # 注册一个简单的计算器工具
    tool_call.register_tool(
        tool_name="calculator",
        description="执行简单的数学计算",
        parameters={
            "type": "object",
            "properties": {
                "operation": {
                    "type": "string",
                    "enum": ["add", "subtract", "multiply", "divide"],
                    "description": "要执行的操作"
                },
                "a": {
                    "type": "number",
                    "description": "第一个数字"
                },
                "b": {
                    "type": "number",
                    "description": "第二个数字"
                }
            },
            "required": ["operation", "a", "b"]
        },
        handler=lambda operation, a, b: {
            "add": a + b,
            "subtract": a - b,
            "multiply": a * b,
            "divide": a / b if b != 0 else "错误: 除数不能为0"
        }.get(operation, "未知操作")
    )
    
    # 发送消息，要求模型使用计算器工具
    response = provider.complete([
        {"role": "user", "content": "计算25乘以16是多少？"}
    ])
    
    print("\n工具调用结果:\n" + response)

def demo_streaming(provider: Any):
    """
    演示流式输出功能
    
    Args:
        provider: 提供商实例
    """
    logger.info("演示流式输出功能")
    
    if not provider.has_capability("streaming"):
        logger.error("当前提供商不支持流式输出功能")
        return
        
    print("\n流式输出:")
    for chunk in provider.complete_stream([
        {"role": "user", "content": "用简洁的语言解释什么是机器学习，不要超过5句话。"}
    ]):
        print(chunk, end="", flush=True)
    print()

def main():
    """主函数"""
    parser = argparse.ArgumentParser(description="Toolkit模式示例")
    parser.add_argument("--api-key", help="API密钥")
    parser.add_argument("--search", action="store_true", help="启用搜索功能")
    args = parser.parse_args()
    
    try:
        # 注册提供商
        register_providers()
        
        # 创建提供商实例
        provider = create_provider("enhanced_moonshot", 
                                  api_key=args.api_key,
                                  search=args.search)
        
        # 获取可用功能
        capabilities = provider.get_available_capabilities()
        logger.info(f"可用功能: {', '.join(capabilities)}")
        
        # 演示功能
        demo_partial_mode(provider)
        
        if args.search:
            demo_web_search(provider)
            
        demo_tool_calls(provider)
        demo_streaming(provider)
        
    except Exception as e:
        logger.error(f"发生错误: {str(e)}")

if __name__ == "__main__":
    main()
